import os

import dotenv
from langchain_classic.chains.conversation.base import ConversationChain
from langchain_classic.chains.llm import LLMChain
from langchain_classic.memory import ConversationBufferMemory, ConversationBufferWindowMemory, \
    ConversationTokenBufferMemory, ConversationSummaryMemory, ConversationSummaryBufferMemory
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate, PromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, RunnableConfig
from langchain_openai import ChatOpenAI

# 会话续写
def fun_1(question):
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )

    # 构建提示词模板
    prompt_template = ChatPromptTemplate.from_messages([
        ("system", "你是一个AI助手")
        , ("human", "{text}")
    ])

    while True:
        chain = prompt_template | chat_model
        response = chain.invoke({"text": question})
        print(f"AI回复：{response.content}")
        #继续获取输入
        user_input = input("请输入你的问题：(输入‘退出’结束会话)")
        #将上面的消息存放到提示词消息列表中
        prompt_template.messages.append(AIMessage(content=response.content))
        prompt_template.messages.append(HumanMessage(content=user_input))
        if user_input == "退出":
            break

#fun_1("你好")

# ChatMessageHistory  使用,一个纯粹的会话历史记录，用于存储会话消息。
from langchain_community.chat_message_histories import ChatMessageHistory
def fun_2():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )
    #创建一个会话历史记录实例
    history = ChatMessageHistory()
    #添加相关消息进行存储
    history.add_user_message("你好，我是张三")
    history.add_ai_message("很高兴认识你，我是李四")
    history.add_user_message("你猜我是谁")
    #打印存储的消息
    print(history.messages)

    response = chat_model.invoke(history.messages)
    print(response.content)

#fun_2()

#ConversationBufferMemory  使用,一个基础的对话记忆组件，专门用于按原始顺序存储完整的对话历史。适用于对话轮次较少，一来完整上下文场景，比如简单的聊天机器人。
#支持2种返回模式 return_messages  =true返回消息对象列表，=fals（默认）e返回文本字符串
def fun_3():
    #创建一个会话历史记录实例
    memory = ConversationBufferMemory()
    #存储相关消息
    memory.save_context({"input": "你好，我是张三"}, {"output": "很高兴认识你，我是李四"})
    memory.save_context({"input": "你猜我是谁"}, {"output": "你是张三啊"})
    #打印存储的消息
    print(memory.load_memory_variables({})) #观察打印的结果，返回的是一个字典，key为history,value为存储的消息字符串

#fun_3()

def fun_4():
    #创建一个会话历史记录实例
    memory = ConversationBufferMemory(return_messages= True)
    #存储相关消息
    memory.save_context({"input": "你好，我是张三"}, {"output": "很高兴认识你，我是李四"})
    memory.save_context({"input": "你猜我是谁"}, {"output": "你是张三啊"})
    #打印存储的消息
    print(memory.load_memory_variables({})) #观察打印的结果，返回的是一个字典，key为history，value为存储的消息对象列表

#fun_4()

#结合 PromptTemplate使用
def fun_5():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )

    #提示词模板,history变量名为如果改成其他的，则不传递此变量值会报错，原因通过上面的例子可以看出，memory返回的key就是 history
    prompt = PromptTemplate.from_template(template="你是一个AI助手。 历史对话：{history} ， 问题：{question} 回复：")
    memory = ConversationBufferMemory() #可以通过参数memory_key=xxx 来修改key,默认为history，修改后上面的history变量也要改成相同的名称
    llm_chain = LLMChain(llm=chat_model, prompt=prompt, memory=memory)
    response = llm_chain.invoke({"question": "你好，我是张三"})
    print(response)
    #继续提问，上次的问题，会自动添加到历史记录中
    response = llm_chain.invoke({"question": "你猜我是谁？"})
    print(response)

#fun_5()

#结合ChatPromptTemplate使用
def fun_6():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )

    #提示词模板,history变量名为如果改成其他的，则不传递此变量值会报错，原因通过上面的例子可以看出，memory返回的key就是 history
    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是一个AI助手")
        , MessagesPlaceholder(variable_name="history") # MessagesPlaceholder消息插入
        , ("human", "问题：{question}")
    ])
    memory = ConversationBufferMemory(return_messages= True) #可以通过参数memory_key=xxx 来修改key,默认为history，修改后上面的history变量也要改成相同的名称
    llm_chain = LLMChain(llm=chat_model, prompt=prompt, memory=memory)
    response = llm_chain.invoke({"question": "你好，我是张三"})
    print(response) #观察打印，发现第一次问题，history里已经有值了，这是因为这种方式为了保障会话结果的一致性，在回复的同时会立即把回复记录到历史记录中
    #继续提问，上次的问题，会自动添加到历史记录中
    response = llm_chain.invoke({"question": "你猜我是谁？"})
    print(response)

#fun_6()

#ConversationChain，一个完整的对话链，包含一个LLMChain和一个Memory组件，用于处理完整的对话。实际上就是对ConversationBufferMemory和LLMChain的封装
def fun_7():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )
    #注意，ConversationChain的输入变量为input，question会报错
    prompt = PromptTemplate.from_template(template="你是一个AI助手。 历史对话：{history} ， 问题：{input} 回复：")

    chain = ConversationChain(llm=chat_model , prompt=prompt)
    response = chain.invoke({"input": "你好，我是张三"})
    print(response)
    # 继续提问，上次的问题，会自动添加到历史记录中
    response = chain.invoke({"input": "你猜我是谁？"})
    print(response)

#fun_7()

#使用ConversationChain默认提供的提示词模板，此模板中的变量名为input和history
def fun_8():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )

    chain = ConversationChain(llm=chat_model)
    response = chain.invoke({"input": "你好，我是张三"})
    print(response)
    # 继续提问，上次的问题，会自动添加到历史记录中
    response = chain.invoke({"input": "你猜我是谁？"})
    print(response)

#fun_8()


#ConversationBufferWindowMemory  用于设置会话历史记录的条数，默认为5。 一般在对话中，太多的历史记录并不会提高模型效果，而且浪费TOKEN。
# 但是对于一些特殊场景，比如对话中的问题回答，可能会有特殊要求，比如只保留最近几条历史记录，那么就可以使用ConversationBufferWindowMemory
def fun_9():
    #创建一个会话历史记录实例
    memory = ConversationBufferWindowMemory(k=2)
    #存储相关消息
    memory.save_context({"input": "你好，我是张三"}, {"output": "很高兴认识你，我是李四"})
    memory.save_context({"input": "你猜我是谁"}, {"output": "你是张三啊"})
    memory.save_context({"input": "国庆节是哪天"}, {"output": "十月一日"})
    #打印存储的消息
    print(memory.load_memory_variables({})) #观察打印的结果，返回的是一个字典，key为history,value为存储的消息字符串

#fun_9()

#ConversationTokenBufferMemory  用于设置会话历史记录的TOKEN数 因为需要计算TOKEN数量，所以在创建实例的时候需要传入大模型
# 如果达到TOKEN限制，会删除最早的历史记录
#注意：有些大模型不支持TOKEN计算，比如阿里百炼的千问3
def fun_10():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )
    #创建一个会话历史记录实例
    memory = ConversationTokenBufferMemory(llm=chat_model , max_token_limit=20)
    #存储相关消息
    memory.save_context({"input": "你好，我是张三"}, {"output": "很高兴认识你，我是李四"})
    memory.save_context({"input": "你猜我是谁"}, {"output": "你是张三啊"})
    memory.save_context({"input": "国庆节是哪天"}, {"output": "十月一日"})
    #打印存储的消息
    print(memory.load_memory_variables({})) #观察打印的结果，返回的是一个字典，key为history,value为存储的消息字符串

#fun_10()


#ConversationSummaryMemory,由大模型自动生成会话摘要，并保存在历史记录中
#例子-1
def fun_11():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )
    #创建一个会话历史记录实例
    memory = ConversationSummaryMemory(llm=chat_model)
    #存储相关消息
    memory.save_context({"input": "你好，我是张三"}, {"output": "很高兴认识你，我是李四"})
    memory.save_context({"input": "我喜欢的颜色是黑色"}, {"output": "黑色偏好者常以低调沉稳示人，不喜张扬，更注重实际行动而非言语表达。他们在压力下保持冷静，擅长理性分析问题，成为团队中可靠的决策者。这种特质源于大脑前额叶皮层的高活跃度，使其在危机中能快速给出解决方案"})
    memory.save_context({"input": "国庆节是哪天"}, {"output": "十月一日"})
    memory.save_context({"input": "我喜欢吃苹果"}, {"output": "苹果富含膳食纤维、维生素C及多酚类物质，具有促进消化、抗氧化等功效，是低热量高营养的健康水果。"})
    #打印存储的消息
    print(memory.load_memory_variables({}))

#fun_11()

#例子-2，在创建ConversationSummaryMemory之前，已经有历史消息，那么可以通过ConversationSummaryMemory.chat_memory.messages属性，将历史消息添加到ConversationSummaryMemory实例中
def fun_12():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )
    history = ChatMessageHistory()
    history.add_user_message("你好，我是张三")
    history.add_ai_message("很高兴认识你，我是李四")
    history.add_user_message("我喜欢的颜色是黑色")
    history.add_ai_message("黑色偏好者常以低调沉稳示人，不喜张扬，更注重实际行动而非言语表达。他们在压力下保持冷静，擅长理性分析问题，成为团队中可靠的决策者。这种特质源于大脑前额叶皮层的高活跃度，使其在危机中能快速给出解决方案")
    history.add_user_message("国庆节是哪天")
    history.add_ai_message("十月一日")
    history.add_user_message("我喜欢吃苹果")
    history.add_ai_message("苹果富含膳食纤维、维生素C及多酚类物质，具有促进消化、抗氧化等功效，是低热量高营养的健康水果。")
    #创建一个会话历史记录实例
    memory = ConversationSummaryMemory.from_messages(llm=chat_model, chat_memory= history)
    print(memory.load_memory_variables({}))

    #观察更新后的输出
    memory.save_context({"human": "我的名字叫赵六"}, {"ai": "你好赵六"})
    print(memory.load_memory_variables({}))

    prompt_with_history = f"基于以下对话历史回答问题:\n{memory.buffer}\n\n用户问题: 你觉得我是一个什么样的人?"

    # 直接调用模型
    response = chat_model.invoke(prompt_with_history)
    print(f"模型回答: {response.content}")

#fun_12()


#ConversationSummaryBufferMemory ,混合记忆
#对超出缓冲区的聊天记忆采用摘要，对最近的N条聊天记忆采用原始消息
#注意：百炼平台的模型不支持TOKEN计算，可以使用GPT测试
def fun_13():
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )
    memory = ConversationSummaryBufferMemory(llm=chat_model, max_token_limit=100 , return_messages=True) #return_messages=True 用消息的方式呈现，max_token_limit=100 设置缓冲区大小
    # 存储相关消息
    memory.save_context({"input": "你好，我是张三"}, {"output": "很高兴认识你，我是李四"})
    memory.save_context({"input": "我喜欢的颜色是黑色"}, {"output": "黑色偏好者常以低调沉稳示人，不喜张扬，更注重实际行动而非言语表达。他们在压力下保持冷静，擅长理性分析问题，成为团队中可靠的决策者。这种特质源于大脑前额叶皮层的高活跃度，使其在危机中能快速给出解决方案"})
    memory.save_context({"input": "国庆节是哪天"}, {"output": "十月一日"})
    memory.save_context({"input": "我喜欢吃苹果"}, {"output": "苹果富含膳食纤维、维生素C及多酚类物质，具有促进消化、抗氧化等功效，是低热量高营养的健康水果。"})
    print(memory.load_memory_variables({}))

fun_13()
#RunnableWithMessageHistory 是 LangChain 中用于管理对话历史的新一代组件，它替代了传统的 ConversationChain 和各种 Memory 组件。
# 主要特点：
# 现代化设计：符合 LangChain 新的 Runnable API 设计理念
# 灵活的历史管理：可以与各种消息历史存储后端集成
# 类型安全：更好地处理 BaseMessage 对象
# 可配置性：支持会话 ID 和多种配置选项
# 核心参数：
# runnable：基础的可运行链
# get_session_history：获取会话历史的函数
# input_messages_key：输入消息的键名
# history_messages_key：历史消息的键名
def fun_0():
    # 初始化模型
    dotenv.load_dotenv()
    chat_model = ChatOpenAI(
        model=os.getenv("MODEL_NAME")
        , base_url=os.getenv("BASE_URL")
        , api_key=os.getenv("API_KEY")
        , temperature=0.8  # 精确度
        , max_tokens=1024  # 回答长度
    )

    # 创建提示词模板
    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是一个AI助手"),
        MessagesPlaceholder(variable_name="history"),
        ("human", "{input}")
    ])

    # 创建基础链
    chain = prompt | chat_model

    # 创建消息历史存储
    store = {}

    def get_session_history(session_id: str):
        if session_id not in store:
            store[session_id] = ChatMessageHistory()
        return store[session_id]

    # 创建带历史记录的链
    chain_with_history = RunnableWithMessageHistory(
        chain,
        get_session_history,
        input_messages_key="input",
        history_messages_key="history"
    )

    # 使用示例
    response = chain_with_history.invoke(
        {"input": "你好，我是张三"},
        config=RunnableConfig(configurable={"session_id": "user1"})  # 修复点
    )
    print(response.content)

    # 继续对话
    response = chain_with_history.invoke(
        {"input": "你猜我是谁？"},
        config=RunnableConfig(configurable={"session_id": "user1"})  # 修复点
    )
    print(response.content)

#fun_0()